package zjl.graph;

import zjl.util.QYPriorityQueue;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class YenTopKShortestPathsAlgorithm {

    private Graph graph;

    // intermediate variables
    private List<Path> resultList = new ArrayList<>();
    private Map<Path, BaseVertex> pathDerivationVertexIndex = new HashMap<>();
    private QYPriorityQueue<Path> pathCandidatesQueue = new QYPriorityQueue<>();

    // the ending vertices of the paths
    private BaseVertex sourceVertex;
    private BaseVertex targetVertex;

    // variables for debugging and testing
    private int generatedPathNum = 0;

    public YenTopKShortestPathsAlgorithm(BaseGraph graph) {
        this(graph, null, null);
    }

    public YenTopKShortestPathsAlgorithm(BaseGraph graph, BaseVertex sourceVertex, BaseVertex targetVertex) {
        if (graph == null) {
            throw new IllegalArgumentException("A null graph object occurs!");
        }

        this.graph = new Graph((Graph) graph);
        this.sourceVertex = sourceVertex;
        this.targetVertex = targetVertex;

        this.init();
    }

    private void init() {
        clear();

        // get the shortest path by default if both source and target exist
        if (this.sourceVertex != null && this.targetVertex != null) {
            Path shortestPath = getShortestPath(sourceVertex, targetVertex);
            if (!shortestPath.getVertexList().isEmpty()) {
                pathCandidatesQueue.offer(shortestPath);
                pathDerivationVertexIndex.put(shortestPath, sourceVertex);
            }
        }
    }

    public void clear() {
        this.pathCandidatesQueue = new QYPriorityQueue<>();
        this.pathDerivationVertexIndex.clear();
        this.resultList.clear();
        this.generatedPathNum = 0;
    }

    public List<Path> getResultList() {
        return this.resultList;
    }

    public int getCandidateSize() {
        return this.pathDerivationVertexIndex.size();
    }

    public int getGeneratedPathNum() {
        return this.generatedPathNum;
    }

    public Path getShortestPath(BaseVertex sourceVertex, BaseVertex targetVertex) {
        DijkstraShortestPathAlgorithm dspa = new DijkstraShortestPathAlgorithm(graph);
        return dspa.getShortestPath(sourceVertex, targetVertex);
    }

    public boolean hasNext() {
        return !this.pathCandidatesQueue.isEmpty();
    }

    public Path next() {
        // 1. prepare for removing vertices and arcs
        Path currPath = this.pathCandidatesQueue.poll();
        this.resultList.add(currPath);

        BaseVertex currDerivation = this.pathDerivationVertexIndex.get(currPath);
        int currPathHash = currPath.getVertexList().subList(0, currPath.getVertexList().indexOf(currDerivation)).hashCode();

        int count = this.resultList.size();

        // 2. remove the vertices and arcs in the graph
        for (int i = 0; i < count - 1; i++) {
            Path currResultPath = resultList.get(i);

            int currDevVertexId = currResultPath.getVertexList().indexOf(currDerivation);

            if (currDevVertexId < 0) {
                continue;
            }

            // Note that the following condition makes sure all candidates should be considered.
            // The algorithm in the paper is not correct for removing some candidates by mistake.
            int pathHash = currResultPath.getVertexList().subList(0, currDevVertexId).hashCode();

            if (pathHash != currPathHash) {
                continue;
            }

            BaseVertex curSuccVertex = currResultPath.getVertexList().get(currDevVertexId + 1);

            graph.deleteEdge(this.graph.getEdgeInstanceOfOriginGraph(currDerivation, curSuccVertex));
        }

        int pathLength = currPath.getVertexList().size();
        List<BaseVertex> currPathVertexList = currPath.getVertexList();
        for (int i = 0; i < pathLength - 1; i++) {
            graph.deleteVertex(currPathVertexList.get(i));
            graph.deleteEdge(this.graph.getEdgeInstanceOfOriginGraph(currPathVertexList.get(i), currPathVertexList.get(i + 1)));
        }

        // 3. calculate the shortest tree rooted at target vertex in the graph
        DijkstraShortestPathAlgorithm reverseTree = new DijkstraShortestPathAlgorithm(graph);
        reverseTree.getShortestPathFlower(targetVertex);

        // 4. recover the deleted vertices and update the cost and identify the new candidate results
        boolean isDone = false;
        for (int i = pathLength - 2; i >= 0 && !isDone; i--) {
            // 4.1 get the vertex to be recovered
            BaseVertex currRecoverVertex = currPathVertexList.get(i);
            this.graph.recoverDeletedVertex(currRecoverVertex);

            // 4.2 check if we should stop continuing in the next iteration
            if (currRecoverVertex.getId() == currDerivation.getId()) {
                isDone = true;
            }

            // 4.3 calculate cost using forward star form
            Path subPath = reverseTree.updateCostForward(currRecoverVertex);

            // 4.4 get one candidate result if possible
            if (subPath != null) {
                generatedPathNum++;

                // 4.4.1 get the prefix from the concerned path
                double cost = 0.0;
                List<BaseVertex> preVertexPath = new ArrayList<>();
                List<BaseEdge> preEdgePath = new ArrayList<>();

                reverseTree.correctCostBackward(currRecoverVertex);

                for (int j = 0; j < pathLength; j++) {
                    BaseVertex currVertex = currPathVertexList.get(j);
                    if (currVertex.getId() == currRecoverVertex.getId()) {
                        j = pathLength;
                    } else {
                        BaseEdge currEdge = this.graph.getEdgeInstanceOfOriginGraph(currPathVertexList.get(j), currPathVertexList.get(j + 1));
                        preEdgePath.add(currEdge);
                        preVertexPath.add(currVertex);
                        cost += currEdge.getWeight();
                    }
                }

                preVertexPath.addAll(subPath.getVertexList());
                preEdgePath.addAll(subPath.getEdgeList());

                // 4.4.2 compose a candidate
                subPath.setWeight(cost + subPath.getWeight());
                subPath.getVertexList().clear();
                subPath.getVertexList().addAll(preVertexPath);
                subPath.getEdgeList().clear();
                subPath.getEdgeList().addAll(preEdgePath);

                // 4.4.3 put it in the candidate pool if new
                if (!this.pathDerivationVertexIndex.containsKey(subPath)) {
                    this.pathCandidatesQueue.offer(subPath);
                    this.pathDerivationVertexIndex.put(subPath, currRecoverVertex);
                }
            }

            // 4.5 restore the edge
            BaseVertex succVertex = currPathVertexList.get(i + 1);
            graph.recoverDeletedEdge(graph.getEdgeInstanceOfOriginGraph(currRecoverVertex, succVertex));

            // 4.6 update cost if necessary
            double tmpcost = graph.getEdgeWeight(currRecoverVertex, succVertex)
                    + reverseTree.getStartVertexDistanceIndex().get(succVertex);

            if (reverseTree.getStartVertexDistanceIndex().get(currRecoverVertex) > tmpcost) {
                reverseTree.getStartVertexDistanceIndex().put(currRecoverVertex, tmpcost);
                reverseTree.getPredecessorIndex().put(currRecoverVertex, succVertex);
                reverseTree.correctCostBackward(currRecoverVertex);
            }
        }

        // 5. restore everything
        graph.recoverDeletedVertices();
        graph.recoverDeletedEdges();

        return currPath;
    }

    public List<Path> getShortestPaths(BaseVertex source, BaseVertex target, int k) {
        sourceVertex = source;
        targetVertex = target;

        this.init();
        int count = 0;
        while (hasNext() && count < k) {
            next();
            ++count;
        }

        return this.resultList;
    }
}
